import torch
from torch import nn, tensor
# from transformers import AutoFeatureExtractor, SegformerForSemanticSegmentation
# import matplotlib.pyplot as plt
from torchvision.transforms import Resize
import torchvision.transforms as T
from torchvision.ops import masks_to_boxes
from torchvision.transforms.functional import crop


# 以下三个函数 get_segmentation、get_masks 和 crop_images 用于分割并裁剪图片中的时尚单品，以供后续使用。
def get_segmentation(extractor, model, image):
    inputs = extractor(images=image, return_tensors="pt")

    outputs = model(**inputs)
    logits = outputs.logits.cpu()

    upsampled_logits = nn.functional.interpolate(
        logits,
        size=image.size[::-1],
        mode="bilinear",
        align_corners=False,
    )

    pred_seg = upsampled_logits.argmax(dim=1)[0]
    return pred_seg


# 返回两个 masks（tensor）列表和 obj_ids（int）
# 来自 Hugging Face 的 mattmdjaga/segformer_b2_clothes 模型
def get_masks(segmentation):
    obj_ids = torch.unique(segmentation)
    obj_ids = obj_ids[1:]
    masks = segmentation == obj_ids[:, None, None]
    return masks, obj_ids


def crop_images(masks, obj_ids, img):
    boxes = masks_to_boxes(masks)
    crop_boxes = []
    for box in boxes:
        crop_box = tensor([box[0], box[1], box[2] - box[0], box[3] - box[1]])
        crop_boxes.append(crop_box)
    preprocess = T.Compose([
        T.Resize(size=(256, 256)),
        T.ToTensor()
    ])
    cropped_images = {}
    for i in range(len(crop_boxes)):
        crop_box = crop_boxes[i]
        cropped = crop(img, crop_box[1].item(), crop_box[0].item(), crop_box[3].item(), crop_box[2].item())
        cropped_images[obj_ids[i].item()] = preprocess(cropped)
    return cropped_images


import os

image_paths = []
for celeb in os.listdir("./photos"):
    for image in os.listdir(f"./photos/{celeb}/"):
        image_paths.append(f"./photos/{celeb}/{image}")

# from milvus import default_server
from pymilvus import utility, connections
from pymilvus import FieldSchema, CollectionSchema, Collection, DataType

# default_server.start()
connections.connect(host="127.0.0.1", port=19530)
DIMENSION = 2048
BATCH_SIZE = 128
COLLECTION_NAME = "fashion"
TOP_K = 3

fields = [
    FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
    FieldSchema(name='filepath', dtype=DataType.VARCHAR, max_length=200),
    FieldSchema(name="name", dtype=DataType.VARCHAR, max_length=200),
    FieldSchema(name="seg_id", dtype=DataType.INT64),
    FieldSchema(name='embedding', dtype=DataType.FLOAT_VECTOR, dim=DIMENSION)
]

schema = CollectionSchema(fields=fields)
collection = Collection(name=COLLECTION_NAME, schema=schema)
index_params = {
    "index_type": "IVF_FLAT",
    "metric_type": "L2",
    "params": {"nlist": 128},
}
collection.create_index(field_name="embedding", index_params=index_params)
collection.load()

# 使用来自 Hugging Face 的 Nvidia ResNet 50 模型生成 embedding 向量。
# 如遇 SSL 证书 URL 错误，请在导入 resnet50 模型前运行此步骤
import ssl

ssl._create_default_https_context = ssl._create_unverified_context
# 并删除最后一层模型输出
embeddings_model = torch.hub.load('NVIDIA/DeepLearningExamples:torchhub', 'nvidia_resnet50', pretrained=True)
embeddings_model = torch.nn.Sequential(*(list(embeddings_model.children())[:-1]))
embeddings_model.eval()


# 将图像转换为向量并插入到 Milvus 向量数据库中
def embed_insert(data, collection, model):
    with torch.no_grad():
        output = model(torch.stack(data[0])).squeeze()
        collection.insert([data[1], data[2], data[3], output.tolist()])


from PIL import Image

data_batch = [[], [], [], []]

for path in image_paths:
    image = Image.open(path)
    path_split = path.split("/")
    name = " ".join(path_split[2].split("_"))
    segmentation = get_segmentation(extractor, model, image)
    masks, ids = get_masks(segmentation)
    cropped_images = crop_images(masks, ids, image)
    for key, image in cropped_images.items():
        data_batch[0].append(image)
        data_batch[1].append(path)
        data_batch[2].append(name)
        data_batch[3].append(key)
        if len(data_batch[0]) % BATCH_SIZE == 0:
            embed_insert(data_batch, collection, embeddings_model)
            data_batch = [[], [], [], []]

if len(data_batch[0]) != 0:
    embed_insert(data_batch, collection, embeddings_model)
collection.flush()


# 使用输入图像查询 Milvus 向量数据库，以检索和上传衣服图像最相似的的前三个结果。
def embed_search_images(data, model):
    with torch.no_grad():
        output = model(torch.stack(data))
    if len(output) > 1:
        return output.squeeze().tolist()
    else:
        return torch.flatten(output, start_dim=1).tolist()


# data_batch[0]是 tensor 列表
# data_batch[1]是图像文件的文件路径（字符串）
# data_batch[2]是图像中人物的名称列表（字符串）
# data_batch[3]是分割键值列表（int）
data_batch = [[], [], [], []]

search_paths = ["./photos/Taylor_Swift/Taylor_Swift_3.jpg", "./photos/Taylor_Swift/Taylor_Swift_8.jpg"]

for path in search_paths:
    image = Image.open(path)
    path_split = path.split("/")
    name = " ".join(path_split[2].split("_"))
    segmentation = get_segmentation(extractor, model, image)
    masks, ids = get_masks(segmentation)
    cropped_images = crop_images(masks, ids, image)
    for key, image in cropped_images.items():
        data_batch[0].append(image)
        data_batch[1].append(path)
        data_batch[2].append(name)
        data_batch[3].append(key)

embeds = embed_search_images(data_batch[0], embeddings_model)
res = collection.search(embeds,
                        anns_field='embedding',
                        param={"metric_type": "L2",
                               "params": {"nprobe": 10}},
                        limit=TOP_K,
                        output_fields=['filepath'])
for index, result in enumerate(res):
    print(index)
    print(result)
